import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim


data = [[34.6, 78.0, 0],
        [30.3, 43.9, 0],
        [35.8, 72.9, 0],
        [60.2, 86.3, 1],
        [79.0, 75.3, 1],
        [45.1, 56.3, 0],
        [61.1, 96.5, 1],
        [75.0, 46.6, 1],
        [76.1, 87.4, 1],
        [84.4, 43.5, 1],
        [95.9, 38.2, 0],
        [75.0, 30.6, 0],
        [82.3, 76.5, 1],
        [69.4, 97.7, 1],
        [39.5, 76.0, 0],
        [54.0, 89.2, 1],
        [69.1, 52.7, 1],
        [67.9, 46.7, 0]]

x0 = list(filter(lambda x: x[-1] == 0., data))
x1 = list(filter(lambda x: x[-1] == 1., data))
plot_x0_0 = [i[0] for i in x0]
plot_x0_1 = [i[1] for i in x0]
plot_x1_0 = [i[0] for i in x1]
plot_x1_1 = [i[1] for i in x1]

plt.plot(plot_x0_0, plot_x0_1, 'ro', label='x_0')
plt.plot(plot_x1_0, plot_x1_1, 'bo', label='x_1')
# plt.show()


class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.fc = nn.Linear(2, 1)
        self.sm = nn.Sigmoid()

    def forward(self, x):
        x = self.fc(x)
        x = self.sm(x)
        return x


logistic_model = LogisticRegression()
criterion = nn.BCELoss()
optimizer = optim.Adam(logistic_model.parameters(), lr=1e-3)

x_data_0 = torch.cat((torch.tensor(plot_x0_0), torch.tensor(plot_x1_0)), dim=0)
x_data_1 = torch.cat((torch.tensor(plot_x0_1), torch.tensor(plot_x1_1)), dim=0)
x_data = torch.stack((x_data_0, x_data_1), dim=1)
y_data = torch.cat((torch.zeros(len(plot_x0_0)), torch.ones(len(plot_x1_0))), dim=0).unsqueeze_(dim=1)


for epoch in range(50000):
    out = logistic_model(x_data)
    loss = criterion(out, y_data)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 1000 == 0:
        print('loss is {}'.format(loss.cpu().item()))

logistic_model.eval()
w0, w1 = logistic_model.fc.weight[0]
bias = logistic_model.fc.bias.cpu().item()
w0 = w0.cpu().item()
w1 = w1.cpu().item()
print(w0, w1, bias)
X = np.array(list(range(30, 100, 1)), dtype=np.float64)
Y = (-w0 * X - bias) / w1
plt.plot(X, Y)

plt.show()
